package id3.controllers;

import id3.models.RowData;
import id3.utils.Constants;
import java.util.ArrayList;

/**
 *
 * @author tungck
 */
public class InformationGain {
    
    public static float calInfoGain(ArrayList<RowData> data, int attrIndex){
        float gain = Entropy.calculateEntropy(data);
        for (String item : Constants.getMapAttrs()[attrIndex].getStringValues()) {
            ArrayList<RowData> subData = calSubData(data, attrIndex, item);
            gain -= subData.size()*Entropy.calculateEntropy(subData)/(float)data.size();
        }
        return gain;
    }
    
    private static ArrayList<RowData> calSubData(ArrayList<RowData> data, int attrIndex, String valueAttr){
        ArrayList<RowData> subData = new ArrayList<>();
        for (RowData row : data) {
            if(row.getRow().get(attrIndex).equals(valueAttr)){
                subData.add(row);
            }
        }
        return subData;
    }

}
